/**
 * Copyright 2023-2023 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef FRAMEWORK_MODEL_UNIFIED_MODEL_H
#define FRAMEWORK_MODEL_UNIFIED_MODEL_H

#include "model_partitions.h"
#include "framework/common/types.h"
#include "framework/model/base_buffer.h"

namespace hiai {
class UnifiedModel {
public:
    UnifiedModel(const uint8_t* modelData, size_t modelSize, bool isPassedFromHidl = false);
    UnifiedModel() = default;
    ~UnifiedModel() = default;

    UnifiedModel(const UnifiedModel& other)
    {
        isValid_ = other.isValid_;
        modelData_ = other.modelData_;
        modelSize_ = other.modelSize_;
        modelHeader_ = other.modelHeader_;
        modelPartitions_ = other.modelPartitions_;
    }

    UnifiedModel& operator=(const UnifiedModel& other)
    {
        if (this == &other) {
            return *this;
        }
        isValid_ = other.isValid_;
        modelData_ = other.modelData_;
        modelSize_ = other.modelSize_;
        modelHeader_ = other.modelHeader_;
        modelPartitions_ = other.modelPartitions_;
        return *this;
    }

    inline const uint8_t* GetModelData() const
    {
        return modelData_;
    }

    inline size_t GetModelSize() const
    {
        return modelSize_;
    }

    inline const Maybe<ModelFileHeader>& GetModelHeader() const
    {
        return modelHeader_;
    }

    ModelType GetModelType() const
    {
        return modelHeader_.Exist() ? static_cast<ModelType>(modelHeader_.Value().modeltype) : IR_API_GRAPH_MODEL;
    }

    uint32_t GetMagic() const
    {
        return modelHeader_.Exist() ? modelHeader_.Value().magic : 0;
    }

    Maybe<ModelPartition> GetModelPartition(ModelPartitionType type, int8_t index = -1) const;
    size_t GetModelPartitionsSize() const;

    inline bool IsValid() const
    {
        return isValid_;
    }

private:
    hiai::Status Init(bool isPassedFromHidl);
    hiai::Status InitModelHeader();

private:
    const uint8_t* modelData_ {nullptr};
    size_t modelSize_ {0};
    Maybe<ModelFileHeader> modelHeader_ {Maybe<ModelFileHeader>(NULL_MAYBE)};
    ModelPartitions modelPartitions_ {};

    bool isValid_ {false};
};
} // namespace hiai

#endif
